package com.cantor.consumer.loadbalance.strategy;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.WeightRandom;
import cn.hutool.core.util.RandomUtil;
import com.cantor.consumer.loadbalance.LoadBalance;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * 加权随机负载均衡策略
 */
public class RandomLoadBalanceStrategy implements LoadBalance {

    @Override
    public <T> T select(Collection<WeightRandom.WeightObj<T>> weightObjs) {
        // 如果集合为空, 返回null
        if (CollUtil.isEmpty(weightObjs)) {
            return null;
        }
        // 否则按权重随机出
        return RandomUtil.weightRandom(weightObjs).next();
    }

    @Override
    public <T> Collection<T> select(Collection<WeightRandom.WeightObj<T>> weightObjs, int nSelected) {
        if(nSelected <= 0){
            return Collections.emptyList();
        }
        List<T> list = new ArrayList<>(nSelected); // 给初始大小, 提升添加速度
        IntStream.range(0,nSelected).forEach(i->{
            list.add(RandomUtil.weightRandom(weightObjs).next());
        });
        return list;
    }
}
